#!/usr/bin/env python3
"""
Enhanced Gripper Point Cloud Utility

This module provides enGprPcd function for enhancing all point cloud types in unifiedData
by adding 8 bounding box points for each gripper/finger related semantic part.
The top 4 points closest to each gripper's center are marked in cyan color.
"""
import numpy as np
from semPcd import semPcd
from scipy.spatial.transform import Rotation
from transPcd import Transform
from termcolor import cprint
from time_it import timeit
class ENGPRPCD:
    def __init__(self, unifiedData: dict, param: str,pcd_keys:str, surpress_print: bool = False, DEBUG_PCD: bool = True, seg_source: bool = False,episode:int=None,ts:int=None):
        self.unifiedData = unifiedData
        self.embodiment = param[0]
        self.target_embodiment=param[1]
        self.show_bbox = (param[2]=='b') if len(param)>=3 else False
        self.surpress_print = surpress_print
        self.DEBUG_PCD = DEBUG_PCD
        self.seg_source = seg_source
        self.pcd_keys=pcd_keys
        self.episode=episode
        self.ts=ts
        self.key_poses = unifiedData.get('key_poses', None)
        self.semantic_indices = semPcd(unifiedData, self.embodiment, surpress_print=True, seg_source=seg_source)
        self.gripper_parts = ['lhand_lfinger','lhand_rfinger','lhand', 'rhand_lfinger', 'rhand_rfinger','rhand']
        self.gripper_indices = {part: self.semantic_indices[part] for part in self.gripper_parts if part in self.semantic_indices}



    def _get_bbox(self, local_points: np.ndarray) -> np.ndarray:
        x_min, y_min, z_min = local_points[:, 0].min(), local_points[:, 1].min(), local_points[:, 2].min()
        x_max, y_max, z_max = local_points[:, 0].max(), local_points[:, 1].max(), local_points[:, 2].max()
        return np.array([[x_min, y_min, z_min], [x_max, y_min, z_min], [x_min, y_max, z_min], [x_max, y_max, z_min],
                        [x_min, y_min, z_max], [x_max, y_min, z_max], [x_min, y_max, z_max], [x_max, y_max, z_max]])

    def _color_closest_4_points(self, bbox_8_points: np.ndarray, ref_point: np.ndarray) -> np.ndarray:
        """Get 4 closest points from 8 bbox points to reference point and mark them cyan."""
        if bbox_8_points.shape[1] == 3:
            bbox_8_points = np.column_stack([bbox_8_points, np.zeros((8, 3))])
        distances = np.linalg.norm(bbox_8_points[:, :3] - ref_point[:3], axis=1)
        closest_indices = np.argsort(distances)[:4]
        bbox_8_points[closest_indices, 3:6] = [0, 255, 255]  # Cyan
        bbox_8_points[np.setdiff1d(np.arange(8), closest_indices), 3:6] = [255, 0, 0]  # Red
        return bbox_8_points

    def _keypose_center(self, part_name: str) -> np.ndarray:
        """Get center point from keypose"""
        side = 'left' if part_name.startswith('lhand') else 'right' if part_name.startswith('rhand') else None
        assert side is not None, f"enGprPcd: Cannot determine side for part {part_name}"
        center_pose = self.key_poses[30:36] if side == 'left' else self.key_poses[12:18]
        return center_pose
    def examine_4points_sameface(self, bbox_local_colored: np.ndarray) -> np.ndarray:
        cyan_mask = np.all(bbox_local_colored[:, 3:6] == [0, 255, 255], axis=1)
        cyan_points = bbox_local_colored[cyan_mask, :3]
        x_coords, y_coords, z_coords = cyan_points[:, 0], cyan_points[:, 1], cyan_points[:, 2]
        if len(np.unique(x_coords)) == 1: return np.array([1, 0, 0]) if x_coords[0] > 0 else np.array([-1, 0, 0])
        elif len(np.unique(y_coords)) == 1: return np.array([0, 1, 0]) if y_coords[0] > 0 else np.array([0, -1, 0])
        elif len(np.unique(z_coords)) == 1: return np.array([0, 0, 1]) if z_coords[0] > 0 else np.array([0, 0, -1])
        # assert False, f"4 cyan points not on same face {cyan_points}"
        cprint(f"4 cyan points not on same face {cyan_points}",'red')
        return np.array([0, 0, 1])

    def _enhance_pcd(self, part_points_local: np.ndarray, norm_vecs: np.ndarray,bbox_local_colored:np.ndarray, part_name: str):
        enhanced_points = Transform.get_transforms(part_points_local, norm_vecs,bbox_local_colored, part_name,stick_faces=True,z_boundary=True,embodiment=self.embodiment,target_embodiment=self.target_embodiment,episode=self.episode,ts=self.ts)
        if self.show_bbox:
            enhanced_points = np.vstack([enhanced_points, bbox_local_colored])
        enhanced_points_global = Transform.local_points_to_global(enhanced_points, self._keypose_center(part_name)[3:6],self._keypose_center(part_name)[:3])
        return enhanced_points_global
    def remove_gripper_points(self,reOrsegPcd:np.ndarray):
        all_gripper_indices = []
        for indices in self.gripper_indices.values():
            all_gripper_indices.extend(indices)
        all_gripper_indices = list(set(all_gripper_indices))
        mask = np.ones(reOrsegPcd.shape[0], dtype=bool)
        mask[all_gripper_indices] = False
        reOrsegPcd = reOrsegPcd[mask]
        return reOrsegPcd
    def get_parts_data(self) -> dict:
        parts_data = {}
        for part, indices in self.gripper_indices.items():
            if not indices: 
                print(f"warning: part {part} has no indices")
                continue
            part_points = self.pcd[indices]
            part_points_local = Transform.global_points_to_local(part_points,self._keypose_center(part)[3:6],self._keypose_center(part)[:3])
            bbox_local = self._get_bbox(part_points_local)
            center_global = self._keypose_center(part)[:3]
            center_local = Transform.global_points_to_local(center_global.reshape(1, 3),self._keypose_center(part)[3:6],self._keypose_center(part)[:3])[0]
            assert np.allclose(center_local, [0,0,0]), f"center_local: {center_local} should be [0,0,0]"
            bbox_local_colored = self._color_closest_4_points(bbox_local, center_local)
            norm_vector = self.examine_4points_sameface(bbox_local_colored)
            parts_data[part] = {'global':{'bbox': Transform.local_points_to_global(bbox_local_colored,self._keypose_center(part)[3:6],self._keypose_center(part)[:3]), 'center': center_global,'part_points': part_points},
                                'local' :{'bbox': bbox_local_colored,                                    'center': center_local, 'part_points': part_points_local,'norm_vector': norm_vector}}
            if not self.surpress_print: 
                print(f"    🔧 ENGRIP | {part}: {len(indices)} points → enhanced")
        return parts_data
    def enhance_parts(self,parts_data:dict):
        hand_norm={}
        enhanced_parts = []
        print(f"parts_data: {parts_data.keys()}")   
        for part, data in parts_data.items():
            if part.endswith('hand'):
                hand_norm[part[0]]=-data['local']['norm_vector'] #OPPOSITE THAN ORIGINAL SO FINGER STICK TO GRIPPER
        print(f"hand_norm: {hand_norm}")   
        for part, data in parts_data.items():
            norm_vecs=[data['local']['norm_vector']]
            if not part.endswith('hand') and part[0] in hand_norm:
                norm_vecs.append(hand_norm[part[0]])
            elif part[0] not in hand_norm:
                print(f"part {part[0]} from {part} not in hand_norm {hand_norm.keys()}")
            enhanced_parts.append(self._enhance_pcd(data['local']['part_points'], norm_vecs, data['local']['bbox'], part))
        
        return np.vstack(enhanced_parts) if enhanced_parts else np.array([])
    def check_re_seg_same_idx(self):    
        assert 'repointcloud' in self.unifiedData and 'segpointcloud' in self.unifiedData and self.unifiedData['repointcloud'] is not None and self.unifiedData['segpointcloud'] is not None, f"repointcloud or segpointcloud is None key:{self.unifiedData.keys()} re:{self.unifiedData.get('repointcloud',None)} seg:{self.unifiedData.get('segpointcloud',None)}"
        assert self.unifiedData['repointcloud'].shape[0] == self.unifiedData['segpointcloud'].shape[0], f"repointcloud ({self.unifiedData['repointcloud'].shape[0]} points) and segpointcloud ({self.unifiedData['segpointcloud'].shape[0]} points) must have the same number of points"
        assert np.allclose(self.unifiedData['repointcloud'][:, :3], self.unifiedData['segpointcloud'][:, :3], atol=1e-6), f"XYZ coordinates of points at same indices in repointcloud and segpointcloud must be identical"
    def add_pcd_to_result(self,pcd:np.ndarray,pcd_key:str):
        if pcd_key in self.unifiedData:
            self.unifiedData[pcd_key]=np.vstack([self.unifiedData[pcd_key],pcd])
            self.unifiedData[pcd_key]=self.remove_gripper_points(self.unifiedData[pcd_key])
        return self.unifiedData
    def enGprPcd(self) -> dict:
        self.check_re_seg_same_idx()
        for key in self.pcd_keys:
            assert key in ["repointcloud","segpointcloud"], f"enGprPcd: {key} is not a valid pcd key.currently only support repointcloud and segpointcloud"
            self.pcd=self.unifiedData[key]
            assert self.pcd.shape[0] > 0, f"enGprPcd: {key} is empty."
            assert self.gripper_indices, "enGprPcd: No gripper/finger parts found."
            parts_data = self.get_parts_data()
            new_pcd=self.enhance_parts(parts_data)
            self.add_pcd_to_result(new_pcd,key)
        # for key in ["state","action"]:
        #     if key in self.unifiedData:
        #         self.unifiedData[key]=Transform.move_state_action_by_js_end(self.unifiedData[key])
        return self.unifiedData
@timeit()
def enGprPcd(unifiedData: dict, param: str,surpress_print: bool = False, DEBUG_PCD: bool = True, seg_source: str = "traverse",episode:int=None,ts:int=None) -> dict:
    engprpcd = ENGPRPCD(unifiedData, param,['repointcloud','segpointcloud'], surpress_print=surpress_print, DEBUG_PCD=DEBUG_PCD, seg_source=seg_source,episode=episode,ts=ts)
    return engprpcd.enGprPcd()


